import pandas as pd
import numpy as np

# ========= 策略描述 ========= #
STRATEGY_DESCRIPTIONS = {
    "breakout": """
        <h3>突破策略 (Breakout)</h3>
        <p><strong>买入信号</strong>：RSI < 70，SMA10 > SMA20（确认趋势）。</p>
        <p><strong>参数</strong>：
            <ul>
                <li>SMA10：10日简单移动平均线</li>
                <li>SMA20：20日简单移动平均线</li>
                <li>RSI：14日相对强弱指数，阈值70</li>
            </ul>
        </p>
    """,
    "macd": """
        <h3>MACD 策略 (Moving Average Convergence Divergence)</h3>
        <p><strong>买入信号</strong>：MACD 线上穿信号线，且 MACD 线 > 0。</p>
        <p><strong>参数</strong>：
            <ul>
                <li>快 EMA：12日指数移动平均线</li>
                <li>慢 EMA：26日指数移动平均线</li>
                <li>信号线：9日 MACD 线指数移动平均线</li>
            </ul>
        </p>
    """
}

# ========= 策略函数 ========= #

# 计算 ATR
def calculate_atr(df, periods=14):
    df["High"] = df["High"].ffill()  # 填充 High 列
    df["Low"] = df["Low"].ffill()    # 填充 Low 列
    df["PrevClose"] = df["Close"].shift(1)  # 预计算前一日收盘价
    df["TR"] = df.apply(
        lambda row: max(
            row["High"] - row["Low"],                     # High - Low
            abs(row["High"] - row["PrevClose"] if not pd.isna(row["PrevClose"]) else row["High"]),
            abs(row["Low"] - row["PrevClose"] if not pd.isna(row["PrevClose"]) else row["Low"])
        ),
        axis=1
    )
    return df["TR"].rolling(window=periods).mean()

# 计算 RSI
def calculate_rsi(df, periods=14):
    """计算 RSI"""
    delta = df["Close"].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=periods).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=periods).mean()
    rs = gain / loss
    return 100 - (100 / (1 + rs))

def strategy_slow_breakout(df: pd.DataFrame) -> pd.DataFrame:
    """突破策略：RSI过滤，SMA确认趋势"""
    # 计算必要指标
    df["SMA5"] = df["Close"].rolling(window=5).mean()
    df["SMA10"] = df["Close"].rolling(window=10).mean()
    df["SMA20"] = df["Close"].rolling(window=20).mean()
    df["High20"] = df["Close"].rolling(20).max()
    df["RSI"] = calculate_rsi(df)
    df["ATR"] = calculate_atr(df)

    df["Signal"] = 0
    buy_condition = (df["RSI"] < 70) & (df["SMA10"] > df["SMA20"]) & (df["SMA10"].shift(1) <= df["SMA20"].shift(1))
    df.loc[buy_condition, "Signal"] = 1
    print(f"[INF] 信号数量: {df['Signal'].sum()}")
    return df

def strategy_fast_breakout(df: pd.DataFrame) -> pd.DataFrame:
    """突破策略：RSI过滤，SMA确认趋势"""
    # 计算必要指标
    df["SMA5"] = df["Close"].rolling(window=5).mean()
    df["SMA10"] = df["Close"].rolling(window=10).mean()
    df["SMA20"] = df["Close"].rolling(window=20).mean()
    df["High20"] = df["Close"].rolling(20).max()
    df["RSI"] = calculate_rsi(df)
    df["ATR"] = calculate_atr(df)

    df["Signal"] = 0
    buy_condition = (df["RSI"] < 70) & (df["SMA5"] > df["SMA10"]) & (df["SMA5"].shift(1) <= df["SMA10"].shift(1))
    df.loc[buy_condition, "Signal"] = 1
    print(f"[INF] 信号数量: {df['Signal'].sum()}")
    return df

def strategy_macd(df: pd.DataFrame) -> pd.DataFrame:
    """MACD 策略：MACD 线与信号线交叉"""
    # 计算必要指标
    df["EMA12"] = df["Close"].ewm(span=12, adjust=False).mean()
    df["EMA26"] = df["Close"].ewm(span=26, adjust=False).mean()
    df["MACD"] = df["EMA12"] - df["EMA26"]
    df["Signal"] = df["MACD"].ewm(span=9, adjust=False).mean()
    df["Signal_Flag"] = 0

    # 买入条件：MACD 上穿信号线，且 MACD > 0
    buy_condition = (df["MACD"] > df["Signal"]) & (df["MACD"].shift(1) <= df["Signal"].shift(1)) & (df["MACD"] > 0)
    df.loc[buy_condition, "Signal_Flag"] = 1
    print(f"[INF] MACD 信号数量: {df['Signal_Flag'].sum()}")
    return df

# ========= 策略映射 ========= #
STRATEGIES = {
    "slow_breakout": strategy_slow_breakout,
    "fast_breakout": strategy_fast_breakout,
}
